import Checkbox from "../../Checkbox";
import classNames from "classnames";
import React, {
  forwardRef,
  useEffect,
  useImperativeHandle,
  useMemo,
  useState,
} from "react";
import {
  flatten,
  isNew,
  ROOT_NODE,
  isRoot,
  getAllChildrenKeys,
  copyRawTreeData,
} from "./utils";
import { TreeNode, TreeNodeArrow } from "./TreeNode";
import type { TreeProps } from "./type";
import { useStyles } from "./styles";
import { isDef, isUndefined } from "@ct0/functions";
import { TreeDataType } from "../type";

export default forwardRef(Tree);

function Tree(
  {
    treeData,
    title = "全部",
    total,
    addChildText = "添加子节点",
    hiddenCount,
    editable,
    checkable,
    disabledCheckAll,
    maxShowDeepth,
    maxCreateDeepth,
    expandAll,
    defaultExpandAll,
    highlightSelected,
    defaultSelected,
    selected: propsSelected,
    checked: propsChecked,
    defaultChecked,
    fullHeight,
    className,
    nodeTitleChangeValidator,
    nodeTitleRender,
    onSelect,
    onChange,
    onNodeTitleChange,
    onNodeRemove,
    onNodeAdd,
    onMockChildAdd,
    onCheck,
    onRootNodeSelect,
  }: TreeProps,
  ref: any
) {
  const classes = useStyles();
  const nodes = flatten({
    treeData,
    verticalLines: [],
    parentKey: ROOT_NODE,
    editable,
    maxShowDeepth,
    maxCreateDeepth,
  });
  const [selectedKey, setSelectedKey] = useState(
    propsSelected || defaultSelected
  );

  useEffect(() => setSelectedKey(propsSelected), [propsSelected]);

  // 找到一个节点的所有父节点
  const getAllParentKeys = (
    tree: TreeDataType,
    key: string | number
  ): Array<string | number> => {
    let stack: Array<string | number> = [];
    let going = true;
    const walker = (tree: TreeDataType, key: number | string) => {
      tree.forEach((item) => {
        if (!going) return;
        else if (item.key === key) going = false;
        else if (item.children) {
          stack.push(item.key);
          walker(item.children, key);
        }
      });
      if (going) stack.pop();
    };
    walker(tree, key);
    return stack;
  };

  const getExpandedKeys = useMemo(() => {
    const expended: Record<string, boolean> = {};
    for (const node of nodes) {
      if (node.isLeaf) continue;
      const key = node.key;
      // 根据当前节点的key所有父节点，并展开
      const parents = selectedKey
        ? getAllParentKeys(treeData, selectedKey)
        : [];
      if (defaultExpandAll || parents.includes(key)) expended[key] = true;
      else expended[key] = false;
    }
    return expended;
  }, [treeData, selectedKey]);

  /**
   * 添加节点时，同时只能添加一个节点，标记当前树中是否存在一个正在添加的新节点，
   * 如果有正在编辑的节点，则禁用添加、编辑、删除全部功能
   *  */
  const [editingKey, setEditingKey] = useState<string | number>();
  // 根节点的展开收起要单独设置，为了保存其它节点的保存和收起的上一次状态
  const [closeAll, setCloseAll] = useState(false);
  const [expandedKeys, setExpand] =
    useState<Record<string, boolean>>(getExpandedKeys);

  useEffect(() => {
    if (defaultExpandAll || expandAll) {
      setExpand(getExpandedKeys);
    }
  }, [getExpandedKeys]);

  const [checkedKeys, setCheckedKeys] = useState(() => new Set(defaultChecked));

  useEffect(() => setCheckedKeys(new Set(propsChecked)), [propsChecked]);

  const isExpanded = (key: string | number) => {
    if (expandAll) return true;
    if (isRoot(key)) return !closeAll;
    return expandedKeys[key] === true;
  };

  const shouldShown = (key: string | number) => {
    if (closeAll) return false;
    return isExpanded(key);
  };

  const clickHandler = (key: string | number) => {
    onSelect?.(key);
    setSelectedKey(key);
  };

  const titleChangeHandler = (key: string | number) => (title: string) => {
    /**
     * 首先判断是否是新增的节点，如果是新增的，那么走创建节点流程
     */
    setEditingKey(undefined);

    if (isNew(key)) {
      onNodeAdd?.(key, title);
      return;
    }

    const newData =
      copyRawTreeData(treeData, (node) => {
        if (node.key == key) return { key: node.key, title };
        else return { key: node.key, title: node.title };
      }) ?? [];

    // 触发单独的更新操作
    onNodeTitleChange?.(key, title, newData);
    onChange?.(newData);
  };

  const expandHandler = (key: string | number, expand: boolean) => {
    /**
     * 展开的时候只设置一个，但是关闭的时候，要把所有子节点都关上
     */
    if (expand) {
      expandedKeys[key] = true;
    } else {
      const keys = getAllChildrenKeys(key, treeData);
      for (const k of keys) expandedKeys[k] = false;
    }
    setExpand({ ...expandedKeys });
  };

  const addMockChildHandler = (key: string | number) => () => {
    onMockChildAdd?.(key);
  };

  useImperativeHandle(ref, () => ({
    setOnEditingKey: (key: string | number) => setEditingKey(key),
  }));

  const checkAll =
    checkedKeys.size === nodes.filter((x) => !x.disableCheck).length;

  const checkAllHandler = () => {
    let keys = nodes.map((i) => i.key);
    if (checkAll) keys = [];
    onCheck?.(keys);
    setCheckedKeys(new Set(keys));
  };

  const checkHandler = (key: string | number, checked: boolean) => {
    if (checked) checkedKeys.add(key);
    else checkedKeys.delete(key);
    const keys = [...checkedKeys.keys()];
    onCheck?.(keys);
    setCheckedKeys(new Set(keys));
  };

  const editing = isDef(editingKey);

  // 展开的节点
  const expandedNodes = useMemo(
    () => nodes.filter((i) => shouldShown(i.parentKey)),
    [nodes, closeAll]
  );

  return (
    <div
      className={classNames(
        classes.tree,
        fullHeight && classes.fullHeight,
        className
      )}
      ref={ref}
    >
      <div className={classes.node}>
        <TreeNodeArrow
          open={!closeAll}
          onClick={() => setCloseAll(!closeAll)}
        />
        {checkable && (
          <Checkbox
            checked={checkAll}
            disabled={disabledCheckAll}
            onChange={checkAllHandler}
          />
        )}
        <span className={classes.treeTitle} onClick={onRootNodeSelect}>
          {title} {hiddenCount ? "" : `(${total ?? nodes.length})`}
        </span>
        <div className={classes.growSpace}></div>
        {editable && !editingKey && (
          <div
            className={classNames(
              classes.hoverButton,
              classes.addChild,
              "add-btn"
            )}
            onClick={addMockChildHandler(ROOT_NODE)}
          >
            {addChildText}
          </div>
        )}
      </div>
      {expandedNodes.map((i) => (
        <TreeNode
          addChildText={addChildText}
          key={i.key}
          currentKey={i.key}
          editable={editingKey ? editingKey === i.key : editable}
          /**
           * 有一种情况下，默认是编辑状态
           * 当节点是刚刚新增的节点，那么默认是输入状态，判断条件是，字段的 key 是不是以 _new_add_ 开头
           */
          isEditing={editingKey === i.key || isNew(i.key)}
          checkable={checkable}
          checked={checkedKeys.has(i.key)}
          selected={i.key == selectedKey}
          title={i.title}
          count={hiddenCount ? undefined : i.count}
          titleRender={nodeTitleRender?.(i)}
          isLast={i.isLast}
          isLeaf={i.isLeaf}
          highlightSelected={highlightSelected}
          open={isExpanded(i.key)}
          verticalLines={i.verticalLines}
          onClick={() => !i.disableClick && clickHandler(i.key)}
          onTitleChange={titleChangeHandler(i.key)}
          nodeTitleChangeValidator={nodeTitleChangeValidator}
          onExpand={(expand) => expandHandler(i.key, expand)}
          disableRemove={editing || i.disableRemove}
          disableEdit={editing || i.disableEdit}
          disableAddChild={editing || i.disableAddChild}
          disableCheck={i.disableCheck}
          onNodeRemove={() => onNodeRemove?.(i.key)}
          onMockChildAdd={addMockChildHandler(i.key)}
          onCancelEdit={() => setEditingKey(undefined)}
          onCheck={(checked) => checkHandler(i.key, checked)}
          onEdit={() => setEditingKey(i.key)}
        />
      ))}
    </div>
  );
}
